import torch
import torch.nn as nn
import torchvision.transforms as transforms
import ssl
import numpy as np
from PIL import Image, ImageDraw, ImageOps
import tkinter as tk
from tkinter import Canvas, Button, Label
import cv2

# 处理SSL证书问题
ssl._create_default_https_context = ssl._create_unverified_context

# 设备配置
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 超参数（需要与训练时保持一致）
input_size = 784
hidden_size = 128
num_classes = 10

# 定义与训练时相同的神经网络模型
class NeuralNet(nn.Module):
    def __init__(self, input_size, hidden_size, num_classes):
        super(NeuralNet, self).__init__()
        self.fc1 = nn.Linear(input_size, hidden_size)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(hidden_size, hidden_size)
        self.fc3 = nn.Linear(hidden_size, num_classes)
        self.dropout = nn.Dropout(0.2)
    
    def forward(self, x):
        out = self.fc1(x)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc2(out)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc3(out)
        return out

# 定义卷积神经网络
class ConvNet(nn.Module):
    def __init__(self):
        super(ConvNet, self).__init__()
        # 第一个卷积层: 输入通道1，输出通道16，卷积核大小5x5
        self.conv1 = nn.Conv2d(1, 16, kernel_size=5, stride=1, padding=2)
        # 最大池化层
        self.pool = nn.MaxPool2d(kernel_size=2, stride=2)
        # 第二个卷积层: 输入通道16，输出通道32，卷积核大小5x5
        self.conv2 = nn.Conv2d(16, 32, kernel_size=5, stride=1, padding=2)
        # 全连接层
        self.fc1 = nn.Linear(32 * 7 * 7, 128)
        self.fc2 = nn.Linear(128, 10)
        self.relu = nn.ReLU()
        self.dropout = nn.Dropout(0.5)
        
    def forward(self, x):
        # 输入形状: (batch_size, 1, 28, 28)
        out = self.conv1(x)           # (batch_size, 16, 28, 28)
        out = self.relu(out)
        out = self.pool(out)          # (batch_size, 16, 14, 14)
        out = self.conv2(out)         # (batch_size, 32, 14, 14)
        out = self.relu(out)
        out = self.pool(out)          # (batch_size, 32, 7, 7)
        out = out.reshape(out.size(0), -1)  # 展平: (batch_size, 32*7*7)
        out = self.fc1(out)           # (batch_size, 128)
        out = self.relu(out)
        out = self.dropout(out)
        out = self.fc2(out)           # (batch_size, 10)
        return out

def load_model(model:str="nn", model_path='mnist_model.pth'):
    """加载训练好的模型"""
    try:
        if model == "nn":
            print("使用全连接神经网络模型")
            model = NeuralNet(input_size, hidden_size, num_classes).to(device)
        else:
            print("使用卷积神经网络模型")
            model = ConvNet().to(device)
        model.load_state_dict(torch.load(model_path, map_location=device))
        model.eval()
        print(f"成功加载模型: {model_path}")
        return model
    except FileNotFoundError:
        print(f"未找到模型文件 '{model_path}'，请先运行训练脚本")
        return None

def predict(model, image_array):
    """对图像数组进行预测"""
    with torch.no_grad():
        # 转换为tensor并进行预处理
        image_tensor = torch.from_numpy(image_array).float()
        
        # 检查模型类型并相应处理输入
        if isinstance(model, ConvNet):
            # 对于卷积网络，需要4D输入: (batch_size, channels, height, width)
            image_tensor = image_tensor.unsqueeze(0).unsqueeze(0).to(device)  # (1, 1, 28, 28)
        else:
            # 对于全连接网络，需要2D输入: (batch_size, features)
            image_tensor = image_tensor.reshape(-1, 28*28).to(device)
        
        # 应用与训练时相同的归一化
        mean = 0.1307
        std = 0.3081
        image_tensor = (image_tensor - mean) / std
        
        # 进行预测
        output = model(image_tensor)
        predicted = torch.argmax(output).item()
        probabilities = torch.nn.functional.softmax(output[0], dim=0)
        confidence = probabilities[predicted].item()
        
        # 获取所有类别的概率
        prob_list = probabilities.cpu().numpy()
        
    return predicted, confidence, prob_list

class DrawApp:
    def __init__(self, root, model):
        self.root = root
        self.root.title("手写数字识别")
        self.root.geometry("600x500")
        
        self.model = model
        self.drawing = False
        self.last_x, self.last_y = 0, 0
        
        
        # 创建结果显示标签
        self.result_label = Label(root, text="请在黑色区域手写一个数字 (0-9)", font=("Arial", 14))
        self.result_label.pack(pady=10)
        
        # 创建概率显示标签
        self.prob_label = Label(root, text="", font=("Arial", 10))
        self.prob_label.pack(pady=5)
        
        # 创建按钮
        button_frame = tk.Frame(root)
        button_frame.pack(pady=10)
        
        self.recognize_button = Button(button_frame, text="识别", command=self.recognize, width=10, height=2)
        self.recognize_button.pack(side=tk.LEFT, padx=5)
        
        self.clear_button = Button(button_frame, text="清空", command=self.clear, width=10, height=2)
        self.clear_button.pack(side=tk.LEFT, padx=5)
        
        # 创建画布
        self.canvas = Canvas(root, width=400, height=400, bg='black')
        self.canvas.pack(pady=20)
        self.canvas.bind("<Button-1>", self.start_drawing)
        self.canvas.bind("<B1-Motion>", self.draw)
        self.canvas.bind("<ButtonRelease-1>", self.stop_drawing)
        # 创建图像对象用于处理手写数字
        self.image = Image.new("L", (400, 400), color=0)  # 灰度图像，黑色背景
        self.draw_img = ImageDraw.Draw(self.image)
        
    def start_drawing(self, event):
        self.drawing = True
        self.last_x, self.last_y = event.x, event.y
        
    def draw(self, event):
        if self.drawing:
            # 在画布上绘制
            self.canvas.create_line((self.last_x, self.last_y, event.x, event.y), 
                                  fill='white', width=15, capstyle=tk.ROUND, smooth=tk.TRUE)
            
            # 在PIL图像上同步绘制
            self.draw_img.line([self.last_x, self.last_y, event.x, event.y], 
                              fill=255, width=15)  # 白色线条
            
            self.last_x, self.last_y = event.x, event.y
            
    def stop_drawing(self, event):
        self.drawing = False
        
    def clear(self):
        self.canvas.delete("all")
        self.image = Image.new("L", (400, 400), color=0)
        self.draw_img = ImageDraw.Draw(self.image)
        self.result_label.config(text="请在黑色区域手写一个数字 (0-9)")
        self.prob_label.config(text="")
        
    def recognize(self):
        if self.model is None:
            self.result_label.config(text="模型未加载")
            return
            
        # 预处理图像
        processed_img = self.preprocess_image()
        
        # 进行预测
        predicted, confidence, probabilities = predict(self.model, processed_img)
        
        # 显示结果
        self.result_label.config(text=f"预测数字: {predicted} (置信度: {confidence:.4f})")
        
        # 显示各类别概率
        prob_text = "各类别概率: "
        for i, prob in enumerate(probabilities):
            prob_text += f"{i}: {prob:.3f}  "
        self.prob_label.config(text=prob_text)
        
    def preprocess_image(self):
        # 将图像转换为28x28大小
        # 首先找到数字的边界框
        bbox = self.image.getbbox()
        if bbox is None:
            # 如果没有绘制任何内容，返回全零数组
            return np.zeros((28, 28))
            
        # 裁剪图像到数字边界
        cropped = self.image.crop(bbox)
        
        # 计算需要填充的大小，保持宽高比
        width, height = cropped.size
        max_dim = max(width, height)
        
        # 创建正方形图像，大小为max_dim
        square_img = Image.new('L', (max_dim, max_dim), 0)
        paste_pos = ((max_dim - width) // 2, (max_dim - height) // 2)
        square_img.paste(cropped, paste_pos)
        
        # 调整为28x28大小
        resized = square_img.resize((20, 20), Image.LANCZOS)
        
        # 创建28x28的图像并在中心放置调整后的图像
        final_img = Image.new('L', (28, 28), 0)
        final_img.paste(resized, (4, 4))
        
        # 转换为numpy数组并归一化到0-1范围
        img_array = np.array(final_img) / 255.0
        
        return img_array

def main():
    # 加载模型
    # model = load_model()
    # model = load_model(model="nn", model_path='mnist_model.pth')
    model = load_model(model="cnn", model_path='mnist_cnn_model.pth')
    if model is None:
        print("请先运行训练脚本生成模型文件")
        return
    
    # 创建GUI应用
    root = tk.Tk()
    app = DrawApp(root, model)
    root.mainloop()

if __name__ == "__main__":
    main()
